从什么都不会到QLoRA分布式DPO（二）

wandb.init(
	project="qwen25-dpo",
	name="qwen25-7b-extraction",
	config={
		"model": model_name,
		"max_length": max_length,
		"lora_r": 32,
		"lora_alpha": 64,
		"learning_rate": 2e-4,
		"epochs": 3,
		"batch_size": 4
	}
)

from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset, Dataset
from trl import DPOConfig, DPOTrainer
from peft import get_peft_model, LoraConfig
import wandb

model_name = "Qwen/Qwen2.5-1.5B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

# wandb 设置
wandb.init(
	project="qwen25-dpo",
	name="qwen25-7b-extraction",
	config={
		"model": model_name,
		"max_length": max_length,
		"lora_r": 8,
		"lora_alpha": 32,
		"learning_rate": 1e-5,
		"epochs": 1,
		"batch_size": 4
	}
)

LoraConfig = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.1,
    bias="none",
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, LoraConfig)

dataset = load_dataset("phimes/DPO-bad-boy-chinese-for-Qwen2.5-extended")
print(dataset)


train_dataset = dataset["train"]
test_dataset = dataset["test"]

training_args = DPOConfig(
    output_dir="./output", 
    num_train_epochs=1,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=4,
    gradient_checkpointing=True,
    learning_rate=1e-5,
    evaluation_strategy="steps", 
    eval_steps=10,
    logging_dir="./logs",
    logging_steps=10,
    # wandb 配置
    report_to="wandb"
)

dpo_trainer = DPOTrainer(
    model=model,
    ref_model=None,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    processing_class=tokenizer,
    args=training_args,
)
  
if __name__ == "__main__":
	try:
	    dpo_trainer.train()
	finally:
		# wandb记录结束
	    wandb.finish()

import torch
import torch.nn.functional as F

# 假设 logits 的形状为 [batch_size, seq_len, vocab_size] 
logits = torch.randn(2, 10, 50257)

logp = F.log_softmax(logits, dim=-1) # 输出形状 [batch_size, seq_len, vocab_size] 

# 提取特定位置（如句末）的 logp 值 
logp_last_token = logp[:, -1, :] # 形状 [batch_size, vocab_size] 
logp_specific_token = logp[:, :, 1234] # 提取词表中第1234个token的logp，形状 [batch_size, seq_len]

def manual_log_softmax(logits: torch.Tensor, dim: int = -1) -> torch.Tensor:
    # 计算最大值（数值稳定）
    max_logits = logits.max(dim=dim, keepdim=True).values
    
    # 减去最大值后计算指数
    stable_logits = logits - max_logits
    
    # 计算对数归一化项
    log_sum_exp = torch.log(torch.sum(torch.exp(stable_logits), dim=dim, keepdim=True))
    
    # 计算 logp
    logp = stable_logits - log_sum_exp
    return logp  

# 验证手动实现与PyTorch内置函数是否一致
logits = torch.randn(2, 10, 50257)
logp_pytorch = F.log_softmax(logits, dim=-1)
logp_manual = manual_log_softmax(logits, dim=-1)

assert torch.allclose(logp_pytorch, logp_manual, atol=1e-6), "手动实现与PyTorch结果不一致！"

from transformers import BitsAndBytesConfig
import torch

bnb_config = BitsAndBytesConfig(
	load_in_4bit=True,
	bnb_4bit_quant_type="nf4",
	bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,  # 自动选择计算类型
	bnb_4bit_use_double_quant=True,
)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-1.5B-Instruct",
    quantization_config=bnb_config
)

def prepare_model_for_kbit_training(model, use_gradient_checkpointing=True, gradient_checkpointing_kwargs=None):
    r"""
    Note this method only works for `transformers` models.

    This method wraps the entire protocol for preparing a model before running a training. This includes:
        1- Cast the layernorm in fp32 2- making output embedding layer require grads 3- Add the upcasting of the lm
        head to fp32 4- Freezing the base model layers to ensure they are not updated during training


    Args:
        model (`transformers.PreTrainedModel`):
            The loaded model from `transformers`
        use_gradient_checkpointing (`bool`, *optional*, defaults to `True`):
            If True, use gradient checkpointing to save memory at the expense of slower backward pass.
        gradient_checkpointing_kwargs (`dict`, *optional*, defaults to `None`):
            Keyword arguments to pass to the gradient checkpointing function, please refer to the documentation of
            `torch.utils.checkpoint.checkpoint` for more details about the arguments that you can pass to that method.
            Note this is only available in the latest transformers versions (> 4.34.1).
    """
    loaded_in_kbit = getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False)
    is_gptq_quantized = getattr(model, "quantization_method", None) == "gptq"
    is_aqlm_quantized = getattr(model, "quantization_method", None) == "aqlm"
    is_eetq_quantized = getattr(model, "quantization_method", None) == "eetq"
    is_torchao_quantized = getattr(model, "quantization_method", None) == "torchao"
    is_hqq_quantized = getattr(model, "quantization_method", None) == "hqq" or getattr(model, "hqq_quantized", False)

    if gradient_checkpointing_kwargs is None:
        gradient_checkpointing_kwargs = {}

    for name, param in model.named_parameters():
        # freeze base model's layers
        param.requires_grad = False

    if (
        not is_gptq_quantized
        and not is_aqlm_quantized
        and not is_eetq_quantized
        and not is_hqq_quantized
        and not is_torchao_quantized
    ):
        # cast all non INT8 parameters to fp32
        for param in model.parameters():
            if (
                (param.dtype == torch.float16) or (param.dtype == torch.bfloat16)
            ) and param.__class__.__name__ != "Params4bit":
                param.data = param.data.to(torch.float32)

    if (
        loaded_in_kbit
        or is_gptq_quantized
        or is_aqlm_quantized
        or is_eetq_quantized
        or is_hqq_quantized
        or is_torchao_quantized
    ) and use_gradient_checkpointing:
        # When having `use_reentrant=False` + gradient_checkpointing, there is no need for this hack
        if "use_reentrant" not in gradient_checkpointing_kwargs or gradient_checkpointing_kwargs["use_reentrant"]:
            # For backward compatibility
            if hasattr(model, "enable_input_require_grads"):
                model.enable_input_require_grads()
            else:

                def make_inputs_require_grad(module, input, output):
                    output.requires_grad_(True)

                model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)

        # To support older transformers versions, check if the model supports gradient_checkpointing_kwargs
        _supports_gc_kwargs = "gradient_checkpointing_kwargs" in list(
            inspect.signature(model.gradient_checkpointing_enable).parameters
        )

        if not _supports_gc_kwargs and len(gradient_checkpointing_kwargs) > 0:
            warnings.warn(
                "gradient_checkpointing_kwargs is not supported in this version of transformers. The passed kwargs will be ignored."
                " if you want to use that feature, please upgrade to the latest version of transformers.",
                FutureWarning,
            )

        gc_enable_kwargs = (
            {} if not _supports_gc_kwargs else {"gradient_checkpointing_kwargs": gradient_checkpointing_kwargs}
        )

        # enable gradient checkpointing for memory efficiency
        model.gradient_checkpointing_enable(**gc_enable_kwargs)
    return model

# 查看可训练参数：
model.print_trainable_parameters()

# 查看所有lora层是否正常
for name, param in model.named_parameters():
    if "lora" in name:
        print(f"{name}: requires_grad={param.requires_grad}")

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from datasets import load_dataset, Dataset
from trl import DPOConfig, DPOTrainer
from peft import get_peft_model, LoraConfig, prepare_model_for_kbit_training
import wandb
import torch
from datetime import datetime

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="fp4",
    bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16,  # 自动选择计算类型
    bnb_4bit_use_double_quant=True,
)

model_name = "Qwen/Qwen2.5-1.5B-Instruct"

# wandb
date_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
wandb.init(project="qwen", name=f"{model_name}_{date_time}")

# 加载模型
model = AutoModelForCausalLM.from_pretrained(model_name,quantization_config=bnb_config, use_cache=False)
try:    
    print(model.model.layers[0].self_attn.q_proj.weight.dtype)
    print(model.model.layers[0].self_attn.q_proj.weight.quant_state)
    print(model.model.layers[0].self_attn.q_proj.weight.quant_type)
    print(type(model.model.layers[0].self_attn.q_proj)) 
except:
    print("no quant")

tokenizer = AutoTokenizer.from_pretrained(model_name)

# 量化模型的工具方法，确保模型是量化后，且配置正确
# model = prepare_model_for_kbit_training(model)

# Lora 设置
lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj", "k_proj", "o_proj", "v_proj", "gate_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
# 应用 Lora
model = get_peft_model(model, lora_config)

model.print_trainable_parameters()
for name, param in model.named_parameters():
    if "lora" in name:
        print(f"{name}: requires_grad={param.requires_grad}")

# 加载数据集
dataset = load_dataset("phimes/DPO-bad-boy-chinese-for-Qwen2.5-extended")

train_dataset = dataset["train"]
test_dataset = dataset["test"]

# 训练配置
training_args = DPOConfig(
    output_dir=f"./output", # 输出目录
    num_train_epochs=1, # 训练轮数，1轮
    per_device_train_batch_size=4, # 训练批次大小，2个
    per_device_eval_batch_size=4, # 评估批次大小，2个
    gradient_accumulation_steps=4, # 梯度累积步数，2步
    gradient_checkpointing=True, # 梯度检查点，True
    learning_rate=1e-5, # 学习率，1e-5
    eval_strategy="steps", # 评估策略, 按步数评估
    eval_steps=10, # 每200步，评估一次
    optim="paged_adamw_8bit",
    logging_dir="./logs", # 日志目录，保存日志
    logging_steps=20, # 每100步，保存一次日志
    bf16=False,
    fp16=True,
    report_to=["wandb"],         # 启用wandb记录，其他还可以tensorboard等工具，但是这里只用wandb。
    run_name=f"{model_name}_{date_time}",     # wandb运行名称，将保存在本地，和init的name可以不一致,
)

dpo_trainer = DPOTrainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    processing_class=tokenizer,
    args=training_args,    
)


if __name__ == "__main__":
    try:
        dpo_trainer.train()
    except Exception as e:
        print(e)
    finally:
        wandb.finish()